import requests
import json
import os
from typing import Union, List

class Chatbot:
    def __init__(self, api_url, model_name):
        self.api_url = api_url
        self.model_name = model_name

    def get_prompt(self, file_path):
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        return content

    def post_http_request(self, messages: List[dict], stop_token: List[str] = []) -> requests.Response:
        headers = {
            "User-Agent": "Test Client",
            "Authorization": "Bearer EMPTY",  # 添加Authorization头
            "Content-Type": "application/json"
        }
        pload = {
            "model": self.model_name,
            "messages": messages,
            "temperature": 0.7,
            "stream": False,
            "stop": stop_token,
            "skip_special_tokens": "false",
            "stop_token_ids": [7]
        }
        response = requests.post(self.api_url, headers=headers, json=pload)
        return response

    def get_response(self, response: requests.Response) -> List[str]:
        data = json.loads(response.content)
        output = [choice["message"]["content"] for choice in data["choices"]]
        return output

    def generate_prompt(self, file_name:Union[str, List[str]], init_prompt:str):
        files = []
        if isinstance(file_name, str):
            files.append(file_name)
        elif isinstance(file_name, list):
            files.extend(file_name)
        for file in files:
            file_path = os.path.join('./split_story', file + '_split.txt')
            target_path = os.path.join('./prompt', file + '.txt')
            messages = [
                {"role": "system", "content": self.get_prompt(init_prompt)},
                {"role": "user", "content": self.get_prompt(file_path)}
            ]

            stop_tokens = ["<|im_end|>", "<|im_start|>", "<|/im_start|>", "<|im_continue|>"]

            response = self.post_http_request(messages, stop_token=stop_tokens)
            output_str = ''.join(self.get_response(response))
            
            with open(target_path, 'w', encoding='utf-8') as f:
                f.write(output_str)
            message = []



# 示例用法
if __name__ == "__main__":
    api_url = "http://localhost:30014/v1/chat/completions"
    model_name = "Yi-1.5-9B-Chat"
    init_prompt = './prompt/QA_model_init_prompt.txt' # 初始化辅助QA模型
    files = ["story1",]
    chatbot = Chatbot(api_url, model_name)
    chatbot.generate_prompt(files, init_prompt)
